import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
import pandas as pd
from pandas.plotting import register_matplotlib_converters
import numpy as np
import tushare as ts

import QUANTAXIS as QA
from alpha import Alphas, delta, sma


register_matplotlib_converters()


# hs300 = ts.get_hs300s()
# hs300.to_csv('hs300.csv', encoding='utf_8_sig')
hs300 = pd.read_csv('hs300.csv')
codes = hs300['code'].to_list()
# codes = ['600496', '300044', '601633', '002415', '000100', '300118', '000333', '000002', '002581']
# codes = ['002581']
data = QA.QA_fetch_stock_day_adv(codes, '2020-01-01', '2021-01-15')
data.data['vwap'] = (data.data['amount'] * 100) / (data.data['volume'] * 10000 + 1)
data.data['normal_vol'] = (data.data['volume'].unstack() / data.data['volume'].unstack().max()).stack()
data.data['returns'] = ((data.data['close'].unstack() - data.data['close'].unstack().shift(1)) / data.data['close'].unstack().shift(1)).stack()
df = data.data
alp = Alphas(data.data)
df['alpha'] = alp.alpha007_5410().stack()
# df[['close', 'normal_vol', 'alpha']].plot()
df0 = df.xs('000100', level='code')
df0['sma_vol'] = sma(df0['volume'], 5)
df0.fillna(0, inplace=True)
# df0[['close', 'normal_vol', 'returns', 'alpha']].plot()
# df0.plot()
figure = plt.figure()
axes1 = figure.add_subplot(411)
axes2 = figure.add_subplot(412)
axes3 = figure.add_subplot(413)
axes4 = figure.add_subplot(414)
axes1.plot(df0.index, df0['close'])
axes2.plot(df0.index, df0[['volume', 'sma_vol']])
axes3.plot(df0.index, df0['alpha'])
axes4.plot(df0.index, delta(df0['close'], 4))
multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)
plt.show()